package main

import (
	"flag"
	"fmt"
	"github.com/gorilla/websocket"
	"log"
	"net"
	"net/http"
)

// 定义一个 WebSocket 升级器
var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	// 允许跨域请求
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// 处理 WebSocket 连接的函数
func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	// 升级 HTTP 连接为 WebSocket 连接
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("WebSocket 升级失败:", err)
		return
	}
	// 确保在函数结束时关闭连接
	defer conn.Close()

	for {
		// 读取客户端发送的消息
		_, message, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("读取消息时发生意外关闭错误: %v", err)
			}
			break
		}
		// 将收到的消息打印到控制台
		fmt.Println("收到日志内容:", string(message))
	}
}

func main2() {
	// 注册 WebSocket 处理函数
	http.HandleFunc("/ssh/agentLogger", handleWebSocket)

	// 定义服务器监听的地址和端口
	addr := "0.0.0.0:8280"
	fmt.Printf("服务器正在监听 %s...\n", addr)
	// 启动 HTTP 服务器
	if err := http.ListenAndServe(addr, nil); err != nil {
		log.Fatal("服务器启动失败:", err)
	}
}
func main() {
	// 定义命令行参数，默认值为空字符串表示使用本机所有IP
	ipPtr := flag.String("ip", "", "服务器监听IP地址（默认使用本机所有可用IP）")
	portPtr := flag.Int("port", 8280, "服务器监听端口（默认8280）")
	flag.Parse()

	// 确定最终使用的IP地址
	var addr string
	if *ipPtr != "" {
		// 使用用户指定的IP
		addr = fmt.Sprintf("%s:%d", *ipPtr, *portPtr)
	} else {
		// 默认使用本机所有可用IP
		addr = fmt.Sprintf("0.0.0.0:%d", *portPtr)
	}

	// 注册WebSocket处理函数
	http.HandleFunc("/ssh/agentLogger", handleWebSocket)

	// 打印服务器监听信息
	fmt.Printf("服务器正在监听 %s...\n", addr)

	// 如果用户没有指定IP，尝试获取并显示实际使用的IP
	if *ipPtr == "" {
		ifaces, err := net.Interfaces()
		if err == nil {
			fmt.Println("可用访问地址:")
			for _, i := range ifaces {
				addrs, _ := i.Addrs()
				for _, addr := range addrs {
					var ip net.IP
					switch v := addr.(type) {
					case *net.IPNet:
						ip = v.IP
					case *net.IPAddr:
						ip = v.IP
					}
					if ip != nil && !ip.IsLoopback() && ip.To4() != nil {
						fmt.Printf("  http://%s:%d/ssh/agentLogger\n", ip.String(), *portPtr)
					}
				}
			}
		}
	}

	// 启动HTTP服务器
	if err := http.ListenAndServe(addr, nil); err != nil {
		log.Fatal("服务器启动失败:", err)
	}
}
